{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "# Actor Critic Method\n",
    "\n",
    "**Author:** [Apoorv Nandan](https://twitter.com/NandanApoorv)<br>\n",
    "**Date created:** 2020/05/13<br>\n",
    "**Last modified:** 2020/05/13<br>\n",
    "**Description:** Implement Actor Critic Method in CartPole environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "This script shows an implementation of Actor Critic method on CartPole-V0 environment.\n",
    "\n",
    "### Actor Critic Method\n",
    "\n",
    "As an agent takes actions and moves through an environment, it learns to map\n",
    "the observed state of the environment to two possible outputs:\n",
    "\n",
    "1. Recommended action: A probabiltiy value for each action in the action space.\n",
    "   The part of the agent responsible for this output is called the **actor**.\n",
    "2. Estimated rewards in the future: Sum of all rewards it expects to receive in the\n",
    "   future. The part of the agent responsible for this output is the **critic**.\n",
    "\n",
    "Agent and Critic learn to perform their tasks, such that the recommended actions\n",
    "from the actor maximize the rewards.\n",
    "\n",
    "### CartPole-V0\n",
    "\n",
    "A pole is attached to a cart placed on a frictionless track. The agent has to apply\n",
    "force to move the cart. It is rewarded for every time step the pole\n",
    "remains upright. The agent, therefore, must learn to keep the pole from falling over.\n",
    "\n",
    "### References\n",
    "\n",
    "- [CartPole](http://www.derongliu.org/adp/adp-cdrom/Barto1983.pdf)\n",
    "- [Actor Critic Method](https://hal.inria.fr/hal-00840470/document)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Setup\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "# Configuration parameters for the whole setup\n",
    "seed = 42\n",
    "gamma = 0.99  # Discount factor for past rewards\n",
    "max_steps_per_episode = 10000\n",
    "env = gym.make(\"CartPole-v0\")  # Create the environment\n",
    "env.seed(seed)\n",
    "eps = np.finfo(np.float32).eps.item()  # Smallest number such that 1.0 + eps != 1.0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Implement Actor Critic network\n",
    "\n",
    "This network learns two functions:\n",
    "\n",
    "1. Actor: This takes as input the state of our environment and returns a\n",
    "probability value for each action in its action space.\n",
    "2. Critic: This takes as input the state of our environment and returns\n",
    "an estimate of total rewards in the future.\n",
    "\n",
    "In our implementation, they share the initial layer.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "num_inputs = 4\n",
    "num_actions = 2\n",
    "num_hidden = 128\n",
    "\n",
    "inputs = layers.Input(shape=(num_inputs,))\n",
    "common = layers.Dense(num_hidden, activation=\"relu\")(inputs)\n",
    "action = layers.Dense(num_actions, activation=\"softmax\")(common)\n",
    "critic = layers.Dense(1)(common)\n",
    "\n",
    "model = keras.Model(inputs=inputs, outputs=[action, critic])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Train\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "optimizer = keras.optimizers.Adam(learning_rate=0.01)\n",
    "huber_loss = keras.losses.Huber()\n",
    "action_probs_history = []\n",
    "critic_value_history = []\n",
    "rewards_history = []\n",
    "running_reward = 0\n",
    "episode_count = 0\n",
    "\n",
    "while True:  # Run until solved\n",
    "    state = env.reset()\n",
    "    episode_reward = 0\n",
    "    with tf.GradientTape() as tape:\n",
    "        for timestep in range(1, max_steps_per_episode):\n",
    "            # env.render(); Adding this line would show the attempts\n",
    "            # of the agent in a pop up window.\n",
    "\n",
    "            state = tf.convert_to_tensor(state)\n",
    "            state = tf.expand_dims(state, 0)\n",
    "\n",
    "            # Predict action probabilities and estimated future rewards\n",
    "            # from environment state\n",
    "            action_probs, critic_value = model(state)\n",
    "            critic_value_history.append(critic_value[0, 0])\n",
    "\n",
    "            # Sample action from action probability distribution\n",
    "            action = np.random.choice(num_actions, p=np.squeeze(action_probs))\n",
    "            action_probs_history.append(tf.math.log(action_probs[0, action]))\n",
    "\n",
    "            # Apply the sampled action in our environment\n",
    "            state, reward, done, _ = env.step(action)\n",
    "            rewards_history.append(reward)\n",
    "            episode_reward += reward\n",
    "\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "        # Update running reward to check condition for solving\n",
    "        running_reward = 0.05 * episode_reward + (1 - 0.05) * running_reward\n",
    "\n",
    "        # Calculate expected value from rewards\n",
    "        # - At each timestep what was the total reward received after that timestep\n",
    "        # - Rewards in the past are discounted by multiplying them with gamma\n",
    "        # - These are the labels for our critic\n",
    "        returns = []\n",
    "        discounted_sum = 0\n",
    "        for r in rewards_history[::-1]:\n",
    "            discounted_sum = r + gamma * discounted_sum\n",
    "            returns.insert(0, discounted_sum)\n",
    "\n",
    "        # Normalize\n",
    "        returns = np.array(returns)\n",
    "        returns = (returns - np.mean(returns)) / (np.std(returns) + eps)\n",
    "        returns = returns.tolist()\n",
    "\n",
    "        # Calculating loss values to update our network\n",
    "        history = zip(action_probs_history, critic_value_history, returns)\n",
    "        actor_losses = []\n",
    "        critic_losses = []\n",
    "        for log_prob, value, ret in history:\n",
    "            # At this point in history, the critic estimated that we would get a\n",
    "            # total reward = `value` in the future. We took an action with log probability\n",
    "            # of `log_prob` and ended up recieving a total reward = `ret`.\n",
    "            # The actor must be updated so that it predicts an action that leads to\n",
    "            # high rewards (compared to critic's estimate) with high probability.\n",
    "            diff = ret - value\n",
    "            actor_losses.append(-log_prob * diff)  # actor loss\n",
    "\n",
    "            # The critic must be updated so that it predicts a better estimate of\n",
    "            # the future rewards.\n",
    "            critic_losses.append(\n",
    "                huber_loss(tf.expand_dims(value, 0), tf.expand_dims(ret, 0))\n",
    "            )\n",
    "\n",
    "        # Backpropagation\n",
    "        loss_value = sum(actor_losses) + sum(critic_losses)\n",
    "        grads = tape.gradient(loss_value, model.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
    "\n",
    "        # Clear the loss and reward history\n",
    "        action_probs_history.clear()\n",
    "        critic_value_history.clear()\n",
    "        rewards_history.clear()\n",
    "\n",
    "    # Log details\n",
    "    episode_count += 1\n",
    "    if episode_count % 10 == 0:\n",
    "        template = \"running reward: {:.2f} at episode {}\"\n",
    "        print(template.format(running_reward, episode_count))\n",
    "\n",
    "    if running_reward > 195:  # Condition to consider the task solved\n",
    "        print(\"Solved at episode {}!\".format(episode_count))\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Visualizations\n",
    "In early stages of training:\n",
    "![Imgur](https://i.imgur.com/5gCs5kH.gif)\n",
    "\n",
    "In later stages of training:\n",
    "![Imgur](https://i.imgur.com/5ziiZUD.gif)\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "actor_critic_cartpole",
   "private_outputs": false,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
